Skip to content

Support broadcasting matmul in col-major layout#81

Merged
ajz34 merged 1 commit into
RESTGroup:masterfrom
ajz34:col-major-matmul-broadcasting
Jun 7, 2026
Merged

Support broadcasting matmul in col-major layout#81
ajz34 merged 1 commit into
RESTGroup:masterfrom
ajz34:col-major-matmul-broadcasting

Conversation

@ajz34

@ajz34 ajz34 commented Jun 7, 2026

Copy link
Copy Markdown
Member

This PR implements broadcasting of matmul in column-major.

This is enhancement, but also considered as behavior change.
Broadcasting matmul is allowed in NumPy and many similar row-major frameworks, but disallowed in Julia and similar col-major tensor frameworks. This PR decides to support NumPy-like matmul broadcasting, but in col-major.

Replace the col-major branch of `LayoutMatMulConfig::layout_matmul`
(which rejected rules 3-7) with a thin wrapper that delegates to the
row-major routine via the identity

    C[t, m, n] = A[t, m, k] @ B[t, k, n]   (row-major)
<=> C[n, m, t] = B[n, k, t] @ A[k, m, t]   (col-major)

i.e. reverse all axes and swap A/B. The returned layouts are reversed
and the A/B fields are swapped back so callers see the original operand
order.

This brings full broadcasting (rules 3-7) to col-major for free,
matching the convention that `DeviceMatMulAPI` impls already use
(see e.g. `device_faer/matmul.rs`, which reverses axes and swaps A/B
before dispatching to the row-major kernel).

Also:
- Update the module-level docs to describe the col-major convention
  (matmul dims are the first two; trailing dims broadcast).
- Replace the test that asserted col-major broadcasting fails with
  positive tests covering rules 3, 4 and 7 (incl. a `1` broadcast).
- Update the device_faer matmul test to exercise col-major broadcasting
  end-to-end via `set_default_order(ColMajor)`.

Co-authored-by: Claude Code <noreply@anthropic.com>
Co-authored-by: glm-5.1 <service@zhipuai.cn>
@ajz34 ajz34 merged commit 0c02572 into RESTGroup:master Jun 7, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant